from envs import REGISTRY as env_REGISTRY
from functools import partial
from components.episode_buffer import EpisodeBatch
import numpy as np
import os
import torch as th

class EpisodeRunner:

    def __init__(self, args, logger):
        self.args = args
        self.logger = logger
        self.batch_size = self.args.batch_size_run
        assert self.batch_size == 1
        self.env = env_REGISTRY[self.args.env](**self.args.env_args)
        if getattr(self.env, "episode_limit", None) is None:
            self.episode_limit = self.env.get_env_info()["episode_limit"]
        else:
            self.episode_limit = self.env.episode_limit
        self.t = 0

        self.t_env = 0

        self.train_returns_original = []
        self.train_returns_additional = []
        self.test_returns_original = []
        self.test_returns_additional = []
        self.train_stats = {}
        self.test_stats = {}

        # Log the first run
        self.log_train_stats_t = -1000000

    def setup(self, scheme, groups, preprocess, mac):
        self.new_batch = partial(EpisodeBatch, scheme, groups, self.batch_size, self.episode_limit + 1,
                                 preprocess=preprocess, device=self.args.device)
        self.mac = mac

    def get_env_info(self):
        return self.env.get_env_info()

    def save_replay(self):
        self.env.save_replay()

    def close_env(self):
        self.env.close()

    def reset(self):
        self.batch = self.new_batch()
        self.env.reset()
        self.t = 0

    def run(self, test_mode=False):
        self.reset()

        terminated = False
        episode_return_original = 0
        episode_return_additional = 0
        self.mac.init_hidden(batch_size=self.batch_size)

        while not terminated:

            pre_transition_data = {
                "state": [self.env.get_state()],
                "avail_actions": [self.env.get_avail_actions()],
                "obs": [self.env.get_obs()]
            }

            self.batch.update(pre_transition_data, ts=self.t)

            # Pass the entire batch of experiences up till now to the agents
            # Receive the actions for each agent at this timestep in a batch of size 1
            actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)

            reward, terminated, env_info = self.env.step(actions[0])
            original_reward = reward
            if 'key' in self.args.env_args and 'mpe' in self.args.env_args['key']:
                env_info['additional_reward'] = env_info['additional_reward'][0]
            if 'additional_reward' in env_info:
                additional_reward = env_info['additional_reward']
            else:
                additional_reward = 0.0
            reward = original_reward + additional_reward * (self.args.t_max - self.t_env) / self.args.t_max
            episode_return_original += original_reward
            episode_return_additional += additional_reward

            post_transition_data = {
                "actions": actions,
                "reward": [(reward,)],
                "terminated": [(terminated != env_info.get("episode_limit", False),)],
            }

            self.batch.update(post_transition_data, ts=self.t)

            self.t += 1

        last_data = {
            "state": [self.env.get_state()],
            "avail_actions": [self.env.get_avail_actions()],
            "obs": [self.env.get_obs()]
        }
        self.batch.update(last_data, ts=self.t)

        # Select actions in the last stored state
        actions = self.mac.select_actions(self.batch, t_ep=self.t, t_env=self.t_env, test_mode=test_mode)
        self.batch.update({"actions": actions}, ts=self.t)

        cur_stats = self.test_stats if test_mode else self.train_stats
        cur_returns_original = self.test_returns_original if test_mode else self.train_returns_original
        cur_returns_additional = self.test_returns_additional if test_mode else self.train_returns_additional
        log_prefix = "test_" if test_mode else ""
        cur_stats.update({k: cur_stats.get(k, 0) + env_info.get(k, 0) for k in set(cur_stats) | set(env_info)})
        cur_stats["n_episodes"] = 1 + cur_stats.get("n_episodes", 0)
        cur_stats["ep_length"] = self.t + cur_stats.get("ep_length", 0)

        if self.args.env == "gymma" and "-4f-" in self.args.env_args["key"]:
            foods = last_data['state'][0][:8].reshape(4,2)
            food = "X"
            if 2 not in (foods == [0,0]).sum(axis=1):
                food = "A"
            if 2 not in (foods == [0,7]).sum(axis=1):
                food = "B"
            if 2 not in (foods == [7,0]).sum(axis=1):
                food = "C"
            if 2 not in (foods == [7,7]).sum(axis=1):
                food = "D"
            cur_stats[f"food_{food}"] = cur_stats.get(f"food_{food}", 0) + 1
        if self.args.env == "gymma" and "SimpleTag" in self.args.env_args["key"]:
            catch = ''.join([str(x) for x in self.env._env.get_dead_prey()])
            cur_stats[f"catch_{catch}"] = cur_stats.get(f"catch_{catch}", 0) + 1
        if self.args.env == "sc2_v2":
            dead_enemy_idx = env_info.get('dead_enemy_idx', -1)
            kill = "X"
            if dead_enemy_idx == 0:
                kill = "A"
            if dead_enemy_idx == 1:
                kill = "B"
            if dead_enemy_idx == 2:
                kill = "C"
            if dead_enemy_idx == 3:
                kill = "D"
            cur_stats[f"kill_{kill}"] = cur_stats.get(f"kill_{kill}", 0) + 1

        if not test_mode:
            self.t_env += self.t

        cur_returns_original.append(episode_return_original)
        cur_returns_additional.append(episode_return_additional)

        if test_mode and (len(self.test_returns_original) == self.args.test_nepisode):
            self._log(cur_returns_original, cur_returns_additional, cur_stats, log_prefix)

            if self.args.env=="gymma" and "-4f-" in self.args.env_args["key"]:
                cur_stats["food_X"] = 0
                cur_stats["food_A"] = 0
                cur_stats["food_B"] = 0
                cur_stats["food_C"] = 0
                cur_stats["food_D"] = 0
            if self.args.env=="gymma" and "SimpleTag" in self.args.env_args["key"]:
                for i in range(5):
                    for j in range(5):
                        cur_stats[f"catch_{i}{j}"] = 0
            if self.args.env == "sc2_v2":
                cur_stats["kill_X"] = 0
                cur_stats["kill_A"] = 0
                cur_stats["kill_B"] = 0
                cur_stats["kill_C"] = 0
                cur_stats["kill_D"] = 0

        elif (not test_mode) and self.t_env - self.log_train_stats_t >= self.args.runner_log_interval:
            self._log(cur_returns_original, cur_returns_additional, cur_stats, log_prefix)
            if hasattr(self.mac.action_selector, "epsilon"):
                self.logger.log_stat("epsilon", self.mac.action_selector.epsilon, self.t_env)
            self.log_train_stats_t = self.t_env

            if self.args.env=="gymma" and "-4f-" in self.args.env_args["key"]:
                cur_stats["food_X"] = 0
                cur_stats["food_A"] = 0
                cur_stats["food_B"] = 0
                cur_stats["food_C"] = 0
                cur_stats["food_D"] = 0
            if self.args.env=="gymma" and "SimpleTag" in self.args.env_args["key"]:
                for i in range(5):
                    for j in range(5):
                        cur_stats[f"catch_{i}{j}"] = 0
            if self.args.env == "sc2_v2":
                cur_stats["kill_X"] = 0
                cur_stats["kill_A"] = 0
                cur_stats["kill_B"] = 0
                cur_stats["kill_C"] = 0
                cur_stats["kill_D"] = 0

        return self.batch

    def _log(self, returns_original, returns_additional, stats, prefix):
        self.logger.log_stat(prefix + "return_original_mean", np.mean(returns_original), self.t_env)
        self.logger.log_stat(prefix + "return_original_std", np.std(returns_original), self.t_env)
        self.logger.log_stat(prefix + "return_additional_mean", np.mean(returns_additional), self.t_env)
        self.logger.log_stat(prefix + "return_additional_std", np.std(returns_additional), self.t_env)
        returns_original.clear()
        returns_additional.clear()

        for k, v in stats.items():
            if k != "n_episodes":
                self.logger.log_stat(prefix + k + "_mean" , v/stats["n_episodes"], self.t_env)
        stats.clear()
